from pymilvus import (
    CollectionSchema,
    FieldSchema,
    DataType,
    Collection, MilvusClient
)
from pymilvus.orm import utility

from src.config import MILVUS_HOST, MILVUS_PORT, MILVUS_DB, MILVUS_USER, MILVUS_PASSWORD
from src.util import ConnectionUtil

client = MilvusClient(
    uri=f"http://{MILVUS_HOST}:{MILVUS_PORT}",
    db_name=MILVUS_DB,
    user=MILVUS_USER,
    password=MILVUS_PASSWORD,
    timeout=80,
)


def create_knowledge_db(collection_name: str):
    # 定义字段
    fields = [
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
        FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=1024),
        FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=768)
    ]
    connection_alais = "temp"
    # 连接到 Milvus
    ConnectionUtil.milvus_connect(connection_alais)
    # 检查集合是否存在
    if utility.has_collection(collection_name, using=connection_alais):
        return f"Collection '{collection_name}' already exists."
    # 创建集合模式
    schema = CollectionSchema(fields=fields, description="知识库")
    collection = Collection(collection_name, schema=schema, using=connection_alais, keywords={"timeout": 80})
    # 创建索引
    index_params = {"index_type": "IVF_FLAT",
                    "metric_type": "L2",
                    "params": {"nlist": 1024}}
    collection.create_index("vector", index_params)
    # 加载集合
    collection.load()
    # 关闭连接
    ConnectionUtil.disconnect_milvus(connection_alais)
    return f"Collection '{collection_name}' created successfully."
